import os
import subprocess
from urllib.parse import urljoin

import m3u8
import requests


class M3u8Downloader:
    def __init__(self):
        self.session = requests.Session()
        self.max_workers = 10
        self.timeout = 30

    def merge_ts_files(self, ts_dir, output_file):
        """合并ts文件为mp4"""
        try:
            # 获取所有ts文件并按序号排序
            ts_files = sorted(
                [f for f in os.listdir(ts_dir) if f.endswith(".ts")],
                key=lambda x: int(x.split("_")[0]),
            )

            # 创建文件列表
            list_file = os.path.join(ts_dir, "filelist.txt")
            with open(list_file, "w", encoding="utf-8") as f:
                for ts_file in ts_files:
                    f.write(f"file '{os.path.join(ts_dir, ts_file)}'\n")

            # 使用ffmpeg合并文件
            cmd = [
                "ffmpeg",
                "-f",
                "concat",
                "-safe",
                "0",
                "-i",
                list_file,
                "-c",
                "copy",
                "-y",
                output_file,
            ]

            subprocess.run(cmd, check=True)

            # 清理临时文件
            os.remove(list_file)
            for ts_file in ts_files:
                os.remove(os.path.join(ts_dir, ts_file))
            os.rmdir(ts_dir)

            print(f"合并完成: {output_file}")
            return True

        except Exception as e:
            print(f"合并文件失败: {str(e)}")
            return False

    def parse_m3u8(self, m3u8_url, m3u8_path, ts_files):
        """解析m3u8文件，获取所有ts文件URL"""
        try:
            # 解析m3u8文件
            m3u8_obj = m3u8.load(m3u8_path)
            if not m3u8_obj.segments:
                print("没有找到可下载的片段")
                return False

            # 获取m3u8基础URL
            base_url = os.path.dirname(m3u8_url) + "/"

            # 获取所有ts文件的URL
            for segment in m3u8_obj.segments:
                # 构建完整的ts文件URL
                ts_url = urljoin(base_url, segment.uri)
                ts_files.add(ts_url)

            return True

        except Exception as e:
            print(f"解析m3u8失败: {str(e)}")
            return False

    def merge_from_cache(self, m3u8_path, output_file, ts_dir):
        """从已保存的ts文件合并为mp4"""
        try:
            # 获取所有ts文件并按序号排序
            ts_files = sorted(
                [f for f in os.listdir(ts_dir) if f.endswith(".ts")],
                key=lambda x: int(x.split(".")[0])
            )

            if not ts_files:
                raise Exception("没有找到ts文件")

            # 创建文件列表
            list_file = os.path.join(ts_dir, "filelist.txt")
            with open(list_file, "w", encoding="utf-8") as f:
                for ts_file in ts_files:
                    f.write(f"file '{os.path.join(ts_dir, ts_file)}'\n")

            # 使用ffmpeg合并文件
            cmd = [
                "ffmpeg",
                "-f",
                "concat",
                "-safe",
                "0",
                "-i",
                list_file,
                "-c",
                "copy",
                "-y",
                output_file,
            ]

            subprocess.run(cmd, check=True)

            # 清理临时文件
            os.remove(list_file)
            for ts_file in ts_files:
                os.remove(os.path.join(ts_dir, ts_file))
            os.rmdir(ts_dir)

            print(f"合并完成: {output_file}")
            return True

        except Exception as e:
            print(f"合并文件失败: {str(e)}")
            return False
